import os
from pathlib import Path
from openai import OpenAI
import json


class ElementAgent:
    # Step 1: Define tools (function descriptions)
    tools = [
        {
            "type": "function",
            "function": {
                "name": "write_to_file",
                "description": "将给定的内容写入指定的文件。",
                "parameters": {
                    "type": "object",
                    "properties": {
                        "filename": {
                            "type": "string",
                            "description": "要写入的文件名",
                        },
                        "content": {
                            "type": "string",
                            "description": "要写入文件的内容"
                        }
                    },
                    "required": ["filename", "content"],
                },
            }
        }
    ]

    def __init__(self, latest_dom_data: str, past_dom_data: str):
        self.client = OpenAI(
            # If environment variable is not configured, replace with Baichuan API Key: api_key="sk-xxx",
            api_key='sk-8dcd9ec46bbe42f3947c71d27a21774d',
            base_url="https://dashscope.aliyuncs.com/compatible-mode/v1",
        )
        self.latest_dom_data = latest_dom_data
        self.past_dom_data = past_dom_data

        with open("login_page_system_prompt.md", "r", encoding="utf-8") as f:
            self.system_prompt = f.read()

    def generate_element_json(self):  # Fixed typo in method name (generte -> generate)
        completion = self.client.chat.completions.create(
            model="qwen3-coder-plus-2025-07-22",
            # Model list: https://help.aliyun.com/zh/model-studio/getting-started/models
            messages=[
                {'role': 'system', 'content': "You are a helpful assistant."},  # Fixed missing closing brace
                {'role': 'user', 'content': f'{self.system_prompt},LATEST_DOM_DATA:{self.latest_dom_data}'}
                # Added self.
            ],
            tool_choice="auto",
            tools=ElementAgent.tools
        )

        # Step 4: Parse model response and call tools
        response = completion.model_dump()
        print("=" * 50)
        print(response)
        model_response_content = None

        if 'choices' in response and len(response['choices']) > 0:
            first_choice = response['choices'][0]
            if 'message' in first_choice:
                model_response_content = first_choice['message']

        # Print model-generated content
        if model_response_content is not None:
            print("=" * 50)
            print("模型返回的结果：")
            print(model_response_content)
        else:
            print("未找到模型返回的内容。")

        if model_response_content and 'tool_calls' in model_response_content:  # Added check for 'tool_calls'
            for tool_call in model_response_content['tool_calls']:
                function_name = tool_call['function']['name']
                arguments = json.loads(tool_call['function']['arguments'])

                if function_name == "write_to_file":
                    # Call tool function
                    result = self.write_to_file(arguments['filename'], arguments['content'])  # Added self.
                    print(f"Tool call result: {result}")

    # Step 2: Implement tool function
    @staticmethod  # Added decorator since it doesn't use self
    def write_to_file(filename: str, content: str):
        """将给定的内容写入指定的文件"""
        try:
            with open(filename, 'w', encoding='utf-8') as f:
                f.write(content)
            return {"status": "success", "message": f"Content written to {filename}"}
        except Exception as e:
            return {"status": "error", "message": str(e)}